# Code for "TSM: Temporal Shift Module for Efficient Video Understanding"
# arXiv:1811.08383
# Ji Lin*, Chuang Gan, Song Han
# {jilin, songhan}@mit.edu, ganchuang@csail.mit.edu


from ops.transforms import *
from torch import nn


class TSN(nn.Module):
    def __init__(self,
                 # num_class,
                 num_segments,
                 modality='RGB',
                 base_model='resnet50',
                 new_length=None,
                 # consensus_type='avg',
                 # before_softmax=True,
                 # dropout=0.8,
                 # img_feature_dim=256,
                 crop_num=1,
                 partial_bn=True,
                 print_spec=True,
                 pretrain='imagenet',
                 is_shift=False,
                 shift_div=8,
                 shift_place='blockres',
                 fc_lr5=False,
                 temporal_pool=False,
                 non_local=False):
        super(TSN, self).__init__()
        self.modality = modality
        self.num_segments = num_segments
        self.reshape = False
        self.crop_num = crop_num
        self.pretrain = pretrain
        # self.before_softma
        # self.dropout = dropout
        # self.consensus_type = consensus_type
        # self.img_feature_dim = img_feature_dim  # the dimension of the CNN feature to represent each frame

        self.is_shift = is_shift
        self.shift_div = shift_div
        self.shift_place = shift_place
        self.base_model_name = base_model
        self.fc_lr5 = fc_lr5
        self.temporal_pool = temporal_pool
        self.non_local = non_local

        # if not before_softmax and consensus_type != 'avg':
        #     raise ValueError("Only avg consensus can be used after Softmax")

        if new_length is None:
            self.new_length = 1 if modality == "RGB" else 5
        else:
            self.new_length = new_length
        if print_spec:
            print(("""
                    Initializing TSN with base model: {}.
                    TSN Configurations:
                        input_modality:     {}
                        num_segments:       {}
                        new_length:         {}
            """.format(base_model, self.modality, self.num_segments, self.new_length)))

        self._prepare_base_model(base_model)

        # feature_dim = self._prepare_tsn(num_class)

        # if self.modality == 'Flow':
        #     print("Converting the ImageNet model to a flow init model")
        #     self.base_model = self._construct_flow_model(self.base_model)
        #     print("Done. Flow model ready...")
        # elif self.modality == 'RGBDiff':
        #     print("Converting the ImageNet model to RGB+Diff init model")
        #     self.base_model = self._construct_diff_model(self.base_model)
        #     print("Done. RGBDiff model ready.")
        assert self.modality == 'RGB'

        # self.consensus = ConsensusModule(consensus_type)

        # if not self.before_softmax:
        #     self.softmax = nn.Softmax()

        self._enable_pbn = partial_bn
        if partial_bn:
            self.partialBN(True)

    # def _prepare_tsn(self, num_class):
    #     feature_dim = getattr(self.base_model, self.base_model.last_layer_name).in_features
    #     if self.dropout == 0:
    #         setattr(self.base_model, self.base_model.last_layer_name, nn.Linear(feature_dim, num_class))
    #         self.new_fc = None
    #     else:
    #         setattr(self.base_model, self.base_model.last_layer_name, nn.Dropout(p=self.dropout))
    #         self.new_fc = nn.Linear(feature_dim, num_class)
    #
    #     std = 0.001
    #     if self.new_fc is None:
    #         normal_(getattr(self.base_model, self.base_model.last_layer_name).weight, 0, std)
    #         constant_(getattr(self.base_model, self.base_model.last_layer_name).bias, 0)
    #     else:
    #         if hasattr(self.new_fc, 'weight'):
    #             normal_(self.new_fc.weight, 0, std)
    #             constant_(self.new_fc.bias, 0)
    #     return feature_dim

    def _prepare_base_model(self, base_model):
        print('=> base model: {}'.format(base_model))

        if 'resnet' in base_model:
            # self.base_model = getattr(torchvision.models, base_model)(True if self.pretrain == 'imagenet' else False)
            self.base_model = getattr(torchvision.models, base_model)(False)
            # self.base_model = resnet50(pretrained=True)
            if self.is_shift:
                print('Adding temporal shift...')
                from ops.temporal_shift import make_temporal_shift
                make_temporal_shift(self.base_model, self.num_segments,
                                    n_div=self.shift_div, place=self.shift_place, temporal_pool=self.temporal_pool)

            if self.non_local:
                raise NotImplementedError
                # print('Adding non-local module...')
                # from ops.non_local import make_non_local
                # make_non_local(self.base_model, self.num_segments)

            # self.base_model.last_layer_name = 'fc'
            # self.input_size = 224
            # self.input_mean = [0.485, 0.456, 0.406]
            # self.input_std = [0.229, 0.224, 0.225]

            self.base_model.avgpool = nn.AdaptiveAvgPool2d(1)
            # self.base_model = torch.nn.Sequential(*list(self.base_model.children())[:-1]) # remove the last fc layer

            # if self.modality == 'Flow':
            #     self.input_mean = [0.5]
            #     self.input_std = [np.mean(self.input_std)]
            # elif self.modality == 'RGBDiff':
            #     self.input_mean = [0.485, 0.456, 0.406] + [0] * 3 * self.new_length
            #     self.input_std = self.input_std + [np.mean(self.input_std) * 2] * 3 * self.new_length
            assert self.modality == 'RGB'
        else:
            raise ValueError('Unknown base model: {}'.format(base_model))

    def train(self, mode=True):
        """
        Override the default train() to freeze the BN parameters
        :return:
        """
        super(TSN, self).train(mode)
        count = 0
        if self._enable_pbn and mode:
            print("Freezing BatchNorm2D except the first one.")
            for m in self.base_model.modules():
                if isinstance(m, nn.BatchNorm2d):
                    count += 1
                    if count >= (2 if self._enable_pbn else 1):
                        m.eval()
                        # shutdown update in frozen mode
                        m.weight.requires_grad = False
                        m.bias.requires_grad = False

    def partialBN(self, enable):
        self._enable_pbn = enable

    def get_optim_policies(self):
        first_conv_weight = []
        first_conv_bias = []
        normal_weight = []
        normal_bias = []
        bn = []

        conv_cnt = 0
        bn_cnt = 0
        for m in self.modules():
            if isinstance(m, torch.nn.Conv2d) or isinstance(m, torch.nn.Conv1d) or isinstance(m, torch.nn.Conv3d):
                ps = list(m.parameters())
                conv_cnt += 1
                if conv_cnt == 1:
                    first_conv_weight.append(ps[0])
                    if len(ps) == 2:
                        first_conv_bias.append(ps[1])
                else:
                    normal_weight.append(ps[0])
                    if len(ps) == 2:
                        normal_bias.append(ps[1])
            elif isinstance(m, torch.nn.BatchNorm2d):
                bn_cnt += 1
                # later BN's are frozen
                if not self._enable_pbn or bn_cnt == 1:
                    bn.extend(list(m.parameters()))
            elif isinstance(m, torch.nn.BatchNorm3d):
                bn_cnt += 1
                # later BN's are frozen
                if not self._enable_pbn or bn_cnt == 1:
                    bn.extend(list(m.parameters()))
            elif len(m._modules) == 0:
                if len(list(m.parameters())) > 0:
                    raise ValueError("New atomic module type: {}. Need to give it a learning policy".format(type(m)))

        return [
            {'params': first_conv_weight, 'lr_mult': 5 if self.modality == 'Flow' else 1, 'decay_mult': 1,
             'name': "first_conv_weight"},
            {'params': first_conv_bias, 'lr_mult': 10 if self.modality == 'Flow' else 2, 'decay_mult': 0,
             'name': "first_conv_bias"},
            {'params': normal_weight, 'lr_mult': 1, 'decay_mult': 1,
             'name': "normal_weight"},
            {'params': normal_bias, 'lr_mult': 2, 'decay_mult': 0,
             'name': "normal_bias"},
            {'params': bn, 'lr_mult': 1, 'decay_mult': 0,
             'name': "BN scale/shift"},
        ]

    def forward(self, input, no_reshape=False):
        if not no_reshape:
            sample_len = (3 if self.modality == "RGB" else 2) * self.new_length # self.new_length = 1 if modality == "RGB" else 5

            if self.modality == 'RGBDiff':
                sample_len = 3 * self.new_length
                input = self._get_diff(input)

            base_out = self.base_model(input.view((-1, sample_len) + input.size()[-2:])) # original input size: n, t*c, h, w; after reshape n*t, c, h, w
        else:
            base_out = self.base_model(input)

        # if self.dropout > 0:
        #     base_out = self.new_fc(base_out)
        #
        # if not self.before_softmax:
        #     base_out = self.softmax(base_out)

        if self.reshape: # self.reshape=True
            if self.is_shift and self.temporal_pool:
                base_out = base_out.view((-1, self.num_segments // 2) + base_out.size()[1:])
            else:
                base_out = base_out.view((-1, self.num_segments) + base_out.size()[1:]) # base_out size: n, t, c, h, w
            # output = self.consensus(base_out)
            # return output.squeeze(1)

        return base_out.squeeze()

    @property
    def feature_dim(self):
        return 2048